package cn.cfanr.myapplication;

import okhttp3.*;
import okio.ByteString;

import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class RKLLMClient {
    private static final String SERVER_URL = "http://192.168.1.211:8080/rkllm_chat";
    private OkHttpClient client = new OkHttpClient();

    public interface OnRKLLMResponseListener {
        void onSuccess(String userMessage, String modelResponse);
        void onFailure(IOException e);
    }

    public void chatWithRKLLM(String userMessage, boolean isStreaming, final OnRKLLMResponseListener listener) {
        new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    // 设置请求头
                    Headers headers = new Headers.Builder()
                            .add("Content-Type", "application/json")
                            .add("Authorization", "not_required")
                            .build();

                    // 准备发送的数据
                    JSONObject data = new JSONObject();
                    try {
                        data.put("model", "your_model_deploy_with_RKLLM_Server");
                        JSONArray messages = new JSONArray();
                        JSONObject message = new JSONObject();
                        message.put("role", "user");
                        message.put("content", userMessage);
                        messages.put(message);
                        data.put("messages", messages);
                        data.put("stream", isStreaming);
                    } catch (JSONException e) {
                        e.printStackTrace();
                        listener.onFailure(new IOException("数据构建失败", e));
                        return;
                    }

                    // 创建请求体
                    RequestBody requestBody = RequestBody.create(MediaType.parse("application/json"), data.toString());

                    // 创建请求
                    Request request = new Request.Builder()
                            .url(SERVER_URL)
                            .headers(headers)
                            .post(requestBody)
                            .build();

                    // 发送请求
                    try (Response response = client.newCall(request).execute()) {
                        if (!response.isSuccessful()) {
                            throw new IOException("Unexpected code " + response);
                        }

                        if (!isStreaming) {
                            // 解析非流式响应
                            String responseText = response.body().string();
                            try {
                                JSONObject jsonResponse = new JSONObject(responseText);
                                JSONArray choices = jsonResponse.getJSONArray("choices");
                                JSONObject message = choices.getJSONObject(choices.length() - 1).getJSONObject("message");
                                String modelResponse = message.getString("content");
                                listener.onSuccess(userMessage, modelResponse);
                            }catch (Exception e)
                            {}
                        } else {
                            // 解析流式响应
                            StringBuilder modelResponse = new StringBuilder();
                            ResponseBody responseBody = response.body();
                            if (responseBody != null) {
                               /*for (ByteString line : responseBody.lines()) {
                                    try {
                                        JSONObject jsonLine = new JSONObject(line.string());
                                        JSONArray choices = jsonLine.getJSONArray("choices");
                                        JSONObject choice = choices.getJSONObject(choices.length() - 1);
                                        if (choice.getString("finish_reason") != "stop") {
                                            JSONObject delta = choice.getJSONObject("delta");
                                            String content = delta.getString("content");
                                            modelResponse.append(content);
                                        }
                                    } catch (JSONException e) {
                                        e.printStackTrace();
                                    }
                                }*/
                                modelResponse.append(responseBody);

                            }
                            listener.onSuccess(userMessage, modelResponse.toString());
                        }
                    }
                } catch (IOException e) {
                    listener.onFailure(e);
                }
            }
        }).start();
    }
}
